import os
import time

from matplotlib import pyplot as plt

from quicksort_v1 import quicksort_v1
from quicksort_v2 import quicksort_v2
from quicksort_v3 import quicksort_v3
from quicksort_v4 import quicksort_v4
from quicksort_v4_1 import quicksort_v4_1
from dataset import generate_dataset
from pickle_utils import save_variable


def brief(sequence, ends_len=5):
    if len(sequence) <= ends_len * 2:
        return str(sequence)
    return str(sequence[:ends_len])[:-1] + ', ..., ' + str(sequence[-ends_len:])[1:]


def save_plot(x, y, filepath):
    fig, ax = plt.subplots()
    ax.plot(x, y, linewidth=2.0)
    ax.set_title('Sort time consumption on datasets with\ndifferent num of repeated elements')
    ax.set_xlabel('num of repeated elements in dataset')
    ax.set_ylabel('consumption (ms)')
    ax.set(xlim=(0, 10 ** 6), xticks=range(0, 10 ** 6 + 1, 10 ** 5),
           ylim=(0, 20 * 1000), yticks=range(0, 20 * 1000 + 1, 1000))

    plt.savefig(filepath)


# Press the green button in the gutter to run the script.
if __name__ == '__main__':
    output_name = 'quicksort_v4'
    func = quicksort_v4
    run_times = 20
    dataset_size = 10 ** 6
    output_path = os.path.join('results', output_name)
    os.makedirs(output_path, exist_ok=True)

    print('output_name:', output_name)
    print('func:', func)

    repeating_nums = [10 ** 6 * i * 10 // 100 for i in range(11)]
    # dataset_idxs = [10]
    consumes = []
    for repeating_num in repeating_nums:
        # filename = f'dataset-{i}.txt'
        print(f'Generating dataset. size is {dataset_size}, repeating num is {repeating_num}')
        datasets = [generate_dataset(dataset_size, repeating_num) for i in range(run_times)]
        print('dataset[0]:', brief(datasets[0]))
        # print(f'{filename}:', brief(dataset))
        start = time.time()
        for i in range(run_times):
            func(datasets[i])
        # dataset.sort()
        elapsed = (time.time() - start) / run_times
        consumes.append(elapsed * 1000)
        print('Sorted[0]:', brief(datasets[0]))
        print(f'Consume: {elapsed * 1000} ms')
        print('=' * 50)

    # repeating_nums = [10 ** 6 * i * 10 // 100 for i in dataset_idxs]
    save_plot(repeating_nums,
              consumes,
              os.path.join(output_path, 'plot.png'))
    save_variable(
        {
            'repeating_nums': repeating_nums,
            'time_consumes': consumes
        },
        os.path.join(output_path, 'time_consumes.pkl')
    )
    print('Finished')
